#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>
#include <xmmintrin.h>
#include "gen_matrix.h"
#include "my_malloc.h"
#include "time.h"

// row major
void mm(double *result, double *a, double *b, int dim_size, int np, int rank, double *temp) {
  int x, y, k;
  double* zz;
  double xx;
  __m128d s1,s2,r1,r;
  for (y = 0; y < dim_size; ++y) {
    if(y%np==rank){
      //broadcast
      zz=b+(dim_size*(y/np));
    }else{
      zz=temp;
    }
    MPI_Bcast(zz, dim_size, MPI_DOUBLE, y%np, MPI_COMM_WORLD);
    for(x=rank;x<dim_size;x+=np){
      r=_mm_sub_pd(r,r);
      xx=0;
      for(k=0; k <dim_size ; k+=2) {
        if(k+1<dim_size){
          s1=_mm_loadu_pd(a+(x/np)*dim_size+k);
          s2=_mm_loadu_pd(zz+k);
          r1=_mm_dp_pd(s1,s2,0x33);
          r=_mm_add_pd(r,r1);
        }else{
          xx= a[(x/np)*dim_size+k] *  zz[k];
        }
      }
      _mm_storel_pd(&result[(x/np)*dim_size+y],r);
      result[(x/np)*dim_size+y]+=xx;
    }
  }
/*
  for (y = 0; y < dim_size; ++y) {
    for (x = 0; x < dim_size; ++x) {
      double r = 0.0;
      for (k = 0; k < dim_size; ++k) {
        r += a[y * dim_size + k] *  b[k * dim_size + x];
      }
      result[y * dim_size + x] = r;
    }
  }
*/
  MPI_Barrier(MPI_COMM_WORLD);
}

void print_matrix(double *result, int dim_size) {
  int x, y;
  for (y = 0; y < dim_size; ++y) {
    for (x = 0; x < dim_size; ++x) {
      printf("%f ", result[y * dim_size + x]);
    }
    printf("\n");
  }
  printf("\n");
}

main(int argc, char *argv[]) {
  int np;
  int rank;
  int ierr;
  MPI_Status status;

  double **r;
  double **result;
  double* temp;
  int i,j,k,z;
  int num_arg_matrices;
  // MPI init
  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD, &np);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);

  if (argc != 4) {
    printf("usage: debug_perf test_set matrix_dimension_size\n");
    exit(1);
  }
  int debug_perf = atoi(argv[1]);
  int test_set = atoi(argv[2]);
  matrix_dimension_size = atoi(argv[3]);
  num_arg_matrices = init_gen_sub_matrix(test_set);

  if(rank==0){
    starttime();
  }
  // allocate arrays
  r = (double **)my_malloc(sizeof(double *) * num_arg_matrices);
  temp= (double*)my_malloc(sizeof(double) * matrix_dimension_size);
  result = (double **)my_malloc(sizeof(double *) * 2);
  result[0] = (double *)my_malloc(sizeof(double) * matrix_dimension_size * (matrix_dimension_size/np));
  result[1] = (double *)my_malloc(sizeof(double) * matrix_dimension_size * (matrix_dimension_size/np));

  // get sub matrices
  for (i = 0; i < num_arg_matrices; ++i) {
    r[i] = (double *)my_malloc(sizeof(double) * matrix_dimension_size * (matrix_dimension_size/np));
    if(i==0){
      if (gen_sub_matrix(rank, test_set, i, r[i], 0, matrix_dimension_size - 1, 1, rank, matrix_dimension_size - 1, np, 1) == NULL) {
        printf("inconsistency in gen_sub_matrix\n");
        exit(1);
      }
    }else{
      if (gen_sub_matrix(rank, test_set, i, r[i], rank, matrix_dimension_size - 1, np, 0, matrix_dimension_size - 1, 1, 0) == NULL) {
        printf("inconsistency in gen_sub_matrix\n");
        exit(1);
      }
    }
  }
  MPI_Barrier(MPI_COMM_WORLD);
  // perform matrix multiplies
  int n = 0;

  mm(result[0], r[0], r[1], matrix_dimension_size, np, rank, temp);
  for (i = 2; i < num_arg_matrices; ++i) {
    mm(result[n ^ 0x1], result[n], r[i], matrix_dimension_size, np, rank, temp);
    n = n ^ 0x1;
  }

  if (debug_perf == 0) {
    // print each of the sub matrices
    if(rank==0){
      printf("argument matrix 0\n");
    }
    MPI_Barrier(MPI_COMM_WORLD);
    //print_matrix(r[i], matrix_dimension_size);
    for(j=0; j<matrix_dimension_size; j++){
      if(j%np==rank){
        for(k=0;k<matrix_dimension_size;k++){
          printf("%lf ",r[0][(j/np)*matrix_dimension_size+k]);
        }
      }
      MPI_Barrier(MPI_COMM_WORLD);
      if(rank==0){
        printf("\n");
      }
      MPI_Barrier(MPI_COMM_WORLD);
    }
    for (i = 1; i < num_arg_matrices; ++i) {
      if(rank==0){
        printf("argument matrix %d\n", i);
      }
      MPI_Barrier(MPI_COMM_WORLD);
      //print_matrix(r[i], matrix_dimension_size);
      for(j=0; j<matrix_dimension_size; j++){
        for(k=0;k<matrix_dimension_size;k++){
          if(k%np==rank){
            printf("%lf ",r[i][(k/np)*matrix_dimension_size+j]);
          }
          MPI_Barrier(MPI_COMM_WORLD);
        }
        if(rank==0){
          printf("\n");
        }
        MPI_Barrier(MPI_COMM_WORLD);
      }
    }
    if(rank==0){
      printf("result matrix\n");
    }
    MPI_Barrier(MPI_COMM_WORLD);
    //print_matrix(result[n], matrix_dimension_size);
    for(j=0; j<matrix_dimension_size; j++){
      if(j%np==rank){
        for(k=0;k<matrix_dimension_size;k++){
            printf("%lf ",result[n][(j/np)*matrix_dimension_size+k]);
          }
      }
      MPI_Barrier(MPI_COMM_WORLD);
      if(rank==0){
        printf("\n");
      }
      MPI_Barrier(MPI_COMM_WORLD);
    }
  } else {
    double sum = 0.0;
    my_free(temp);
    double rec;
    //for (i = 0; i < matrix_dimension_size * matrix_dimension_size; ++i) {
      //sum += result[n][i];
      //sum++;
    //}

    for(i=0;i<matrix_dimension_size;i++){
      if(i%np==rank){
        for(j=0;j<matrix_dimension_size;j++){
          sum+=result[n][(i/np)*matrix_dimension_size+j];
        }
      }
    }
    MPI_Reduce(&sum,&rec,1,MPI_DOUBLE,MPI_SUM,0,MPI_COMM_WORLD);
    if(rank==0){ 
      printf("%lf\n", rec);
    }
  }
  MPI_Barrier(MPI_COMM_WORLD);
  if(rank==0){
    endtime();
    printf("time = %d\n", (int)timelapse());
  }
  MPI_Finalize();
}

